package com.renne.ai.dbcmcp.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.renne.ai.dbcmcp.service.DBTableStructureService;
import com.renne.ai.dbcmcp.service.RolePromptService;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.StdioServerTransportProvider;
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpServerTransportProvider;
import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.reactive.function.server.RouterFunction;

import javax.annotation.Resource;
import java.util.List;

@Configuration
public class McpServerConfig {

    @Resource
    private DBTableStructureService dbTableStructureService;
    @Resource
    private RolePromptService rolePromptService;

    // STDIO transport
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "stdio")
    public StdioServerTransportProvider stdioServerTransportProvider() {
        return new StdioServerTransportProvider();
    }

    // SSE transport
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
    public WebFluxSseServerTransportProvider sseServerTransportProvider() {
        return new WebFluxSseServerTransportProvider(new ObjectMapper(), "/mcp/message");
    }

    // Router function for SSE transport used by Spring WebFlux to start an HTTP
    // server.
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
    public RouterFunction<?> mcpRouterFunction(WebFluxSseServerTransportProvider transportProvider) {
        return transportProvider.getRouterFunction();
    }


    @Bean
    public McpSyncServer mcpServer(McpServerTransportProvider transportProvider) {

        // Configure server capabilities with resource support
        var capabilities = McpSchema.ServerCapabilities.builder()
                .tools(true) // Tool support with list changes notifications
                .logging() // Logging support
                .build();

        // 从 dbTableStructureService 中提取 ToolCallback
        ToolCallback[] toolCallbacks = ToolCallbacks.from(dbTableStructureService, rolePromptService);

        // 将 ToolCallback 转换为 SyncToolSpecification
        List<McpServerFeatures.SyncToolSpecification> syncToolSpecifications = McpToolUtils.toSyncToolSpecifications(toolCallbacks);

        // Create the server with both tool and resource capabilities
        McpSyncServer server = McpServer.sync(transportProvider)
                .serverInfo("DBC MCP Server", "0.0.1")
                .capabilities(capabilities)
                .tools(syncToolSpecifications) // 使用转换后的 toolSpecifications
                .build();

        return server;
    }


}
